import bpy
import json
import os
import math
import mathutils

# ===============================================================
# CONFIGURATION - 请在此处修改参数
# ===============================================================

# 保存图片和JSON文件的输出文件夹路径
# Windows示例: "C:/blender_renders/my_scene"
# Linux/macOS示例: "/home/user/blender_renders/my_scene"
output_dir = "C:/blender_renders/alpha_output"

# 环绕拍摄的半径
radius = 10.0

# 每个环上的照片数量
num_images_per_ring = 16

# 俯仰角度 (上: 45度, 中: 0度, 下: -45度)
elevation_angles_deg = [45, 0, -45]

# 相机的垂直视场角 (Field of View in Y-axis)
fov_y_degrees = 50.0

# 渲染分辨率
render_width = 1024
render_height = 1024

# ===============================================================
# SCRIPT LOGIC - 请勿修改以下内容
# ===============================================================

def setup_compositor_for_alpha_output():
    """配置合成器节点以输出Alpha通道到所有颜色通道"""
    scene = bpy.context.scene
    scene.use_nodes = True
    tree = scene.node_tree

    # 清理现有节点
    for node in tree.nodes:
        tree.nodes.remove(node)

    # 创建必要的节点
    render_layers = tree.nodes.new('CompositorNodeRLayers')
    composite_node = tree.nodes.new('CompositorNodeComposite')
    
    # 创建分离和合并RGBA的节点
    sep_rgba = tree.nodes.new('CompositorNodeSepRGBA')
    comb_rgba = tree.nodes.new('CompositorNodeCombRGBA')

    # 将Alpha通道连接到R, G, B, A
    tree.links.new(render_layers.outputs['Image'], sep_rgba.inputs['Image'])
    tree.links.new(sep_rgba.outputs['Alpha'], comb_rgba.inputs['R'])
    tree.links.new(sep_rgba.outputs['Alpha'], comb_rgba.inputs['G'])
    tree.links.new(sep_rgba.outputs['Alpha'], comb_rgba.inputs['B'])
    tree.links.new(sep_rgba.outputs['Alpha'], comb_rgba.inputs['A'])
    
    # 将最终结果连接到输出
    tree.links.new(comb_rgba.outputs['Image'], composite_node.inputs['Image'])
    print("Compositor setup complete for alpha-only output.")

def main():
    """主执行函数"""
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"Created output directory: {output_dir}")

    # 设置渲染引擎和输出格式
    scene = bpy.context.scene
    scene.render.engine = 'CYCLES'  # EEVEE或CYCLES都可以，CYCLES处理透明度通常更精确
    scene.render.film_transparent = True  # 必须开启，以获得有效的Alpha通道
    scene.render.image_settings.file_format = 'PNG'
    scene.render.image_settings.color_mode = 'RGBA'
    scene.render.resolution_x = render_width
    scene.render.resolution_y = render_height
    
    # 设置合成器
    setup_compositor_for_alpha_output()

    # 获取或创建相机
    cam_obj = scene.camera
    if not cam_obj:
        cam_data = bpy.data.cameras.new('PinholeCamera')
        cam_obj = bpy.data.objects.new('PinholeCamera', cam_data)
        scene.collection.objects.link(cam_obj)
        scene.camera = cam_obj
    
    # 设置相机参数
    cam_obj.data.type = 'PERSP'
    cam_obj.data.angle_y = math.radians(fov_y_degrees)

    # 存储所有相机机位信息
    all_camera_data = []
    
    # 定义相机朝向的目标点和上方向
    look_at_point = mathutils.Vector((0.0, 0.0, 0.0))
    up_vector = mathutils.Vector((0.0, 0.0, 1.0))
    
    img_index = 0
    total_images = len(elevation_angles_deg) * num_images_per_ring

    # 开始循环拍摄
    for elevation_deg in elevation_angles_deg:
        elevation_rad = math.radians(elevation_deg)
        
        for i in range(num_images_per_ring):
            azimuth_rad = (i / num_images_per_ring) * 2 * math.pi
            
            # 计算相机位置 (球坐标转笛卡尔坐标)
            x = radius * math.cos(azimuth_rad) * math.cos(elevation_rad)
            y = radius * math.sin(azimuth_rad) * math.cos(elevation_rad)
            z = radius * math.sin(elevation_rad)
            cam_location = mathutils.Vector((x, y, z))
            
            # 设置相机位置
            cam_obj.location = cam_location
            
            # 设置相机朝向
            direction = look_at_point - cam_location
            rot_quat = direction.to_track_quat('-Z', 'Y')
            cam_obj.rotation_euler = rot_quat.to_euler()

            # 渲染当前视图
            print(f"Rendering image {img_index + 1}/{total_images}...")
            filepath = os.path.join(output_dir, f"render_{img_index:03d}.png")
            scene.render.filepath = filepath
            bpy.ops.render.render(write_still=True)
            
            # 保存相机数据
            shot_data = {
                "image_file": os.path.basename(filepath),
                "position": [cam_location.x, cam_location.y, cam_location.z],
                "look_at": [look_at_point.x, look_at_point.y, look_at_point.z],
                "up": [up_vector.x, up_vector.y, up_vector.z],
                "fov_y": fov_y_degrees
            }
            all_camera_data.append(shot_data)
            
            img_index += 1

    # 将相机数据写入JSON文件
    json_path = os.path.join(output_dir, "camera_data.json")
    with open(json_path, 'w') as f:
        json.dump(all_camera_data, f, indent=4)
        
    print(f"\nAll {total_images} images rendered.")
    print(f"Camera data saved to: {json_path}")
    print("Script finished successfully.")

# --- 执行脚本 ---
if __name__ == "__main__":
    main()